import glob
import os
import time
from datetime import datetime

import gym
import numpy as np
import roboschool
import torch

from PPO import PPO


################################### Training ###################################
def train():
    print("============================================================================================")

    ####### initialize environment hyperparameters ######
    # env_name:str -> RoboschoolWalker2d-v1 是使用的 gym 封装过的 roboschool 环境
    """
    状态空间:
        在 gym 官网 https://www.gymlibrary.dev/environments/mujoco/walker2d/ 给出的是 17维
        1. z-coordinate of the top                             顶部的 z 坐标
        2. angle of the top                                    顶部的角度
        3. angle of the thigh joint                            右髋关节的角度
        4. angle of the leg joint                              右腿关节的角度
        5. angle of the foot joint                             右脚关节的角度
        6. angle of the left thigh joint                       左髋关节的角度
        7. angle of the left leg joint                         左腿关节的角度
        8. angle of the left foot joint                        左脚关节的角度
        9. velocity of the x-coordinate of the top             顶部的 x 坐标的速度
        10. velocity of the z-coordinate (height) of the top   顶部的 z 坐标的速度
        11. angular velocity of the angle of the top           顶部的角速度
        12. angular velocity of the thigh hinge                右髋关节的角速度
        13. angular velocity of the leg hinge                  右腿关节的角速度
        14. angular velocity of the foot hinge                 右脚关节的角速度
        15. angular velocity of the thigh hinge                左髋关节的角速度
        16. angular velocity of the leg hinge                  左腿关节的角速度
        17. angular velocity of the foot hinge                 左脚关节的角速度

        但是这里是 22 维，看源码没有看出为什么
        推测是这样的
            （1）关节角度（Joint Angles）
            机器人各个关节的当前角度。
            RoboschoolWalker2d-v1 中的两足机器人通常有 6 个关节（例如髋关节、膝关节和踝关节，每条腿 3 个关节）。
            维度：6 维。

            （2）关节角速度（Joint Angular Velocities）
            机器人各个关节的当前角速度（即关节旋转的速度）。
            维度：6 维。

            （3）躯干位置（Torso Position）
            机器人躯干在三维空间中的位置（x, y, z）。
            维度：3 维。

            （4）躯干线速度（Torso Linear Velocity）
            机器人躯干在三维空间中的线速度（x, y, z 方向的速度）。
            维度：3 维。

            （6）躯干方向（Torso Orientation）
            机器人躯干的朝向，通常用四元数或欧拉角表示。
            维度：3 维（如果使用欧拉角）或 4 维（如果使用四元数）。在 RoboschoolWalker2d-v1 中，通常使用 3 维的欧拉角。

            （6）躯干角速度（Torso Angular Velocity）
            机器人躯干的旋转速度（绕 x, y, z 轴的角速度）。
            维度：3 维。
    
    动作空间：6 个电机的扭矩
        1. Torque applied on the thigh rotor
        2. Torque applied on the leg rotor
        3. Torque applied on the foot rotor
        4. Torque applied on the left thigh rotor
        5. Torque applied on the left leg rotor
        6. Torque applied on the left foot rotor
        
    reward：
        奖励函数是基于以下的奖励：
        1. healthy_reward：如果没有倒地，就返回固定的正奖励
        2. forward_reward：x 方向的速度和加速度
        3. ctrl_cost：能量消耗，扭矩的平方
        
        三个 reward 通过系数加权
    
    终止条件：
        or 关系
        1. 状态空间的任何一个值是无限的(超过了浮点数表达的上下限、仿真器不收敛)
        2. 躯干的顶点低于某个高度
        3. 超过了仿真最大时间
        4. 一些关节的转角超过了某范围
        
    起始条件：
        1. 各个关节的起始角度
    
    """
    env_name = "RoboschoolWalker2d-v1"

    # has_continuous_action_space:bool -> 是否是连续动作空间，动作空间一共有 3 种，分别是离散动作空间、连续动作空间、混合动作空间
    has_continuous_action_space = True  # continuous action space; else discrete

    # max_ep_len:int -> 在一次 episode 中，仿真的最大的步长
    max_ep_len = 1000                   # max timesteps in one episode
    
    # max_training_timesteps:int -> 最大的训练时间步数，3e6 == 3000000 个步长
    max_training_timesteps = int(3e6)   # break training loop if timeteps > max_training_timesteps

    # print_freq:int -> 多少个时间步长才打印该时间步长内的平均奖励
    print_freq = max_ep_len * 10        # print avg reward in the interval (in num timesteps)
    # log_freq:int -> 多少个时间步长就记录一次平均奖励
    log_freq = max_ep_len * 2           # log avg reward in the interval (in num timesteps)
    # save_mode_freq:int -> 多少个时间步长后回保存一下模型的参数
    save_model_freq = int(1e5)          # save model frequency (in num timesteps)

    # action_std:float -> 动作分布(多元正态分布)的初始标准差
    action_std = 0.6                    # starting std for action distribution (Multivariate Normal)
    # action_std_decay_rate:float -> 动作分布的标准差线性衰减率
    action_std_decay_rate = 0.05        # linearly decay action_std (action_std = action_std - action_std_decay_rate)
    # min_action_std:float -> 动作分布的最小标准差，标准差约大，数据约分散
    min_action_std = 0.1                # minimum action_std (stop decay after action_std <= min_action_std)
    # action_std_decay_freq:int -> 动作分布标准差衰减频率，每多少个时间步长就衰减一次，2.5e5 == 250000 个时间步长
    action_std_decay_freq = int(2.5e5)  # action_std decay frequency (in num timesteps)
    
    # 源码阅读疑问：
    # Q1: print_freq 的时间步长的单位是指什么？是指每个 env.step() 一次吗？
    
    #####################################################

    ## Note : print/log frequencies should be > than max_ep_len

    ################ PPO hyperparameters ################
    # PPO 超参数
    # update_timestep:int -> 每多少个时间步长更新一次 policy
    update_timestep = max_ep_len * 4      # update policy every n timesteps
    # K_epochs:int -> 多少 epochs 后更新 PPO policy 的参数
    K_epochs = 80               # update policy for K epochs in one PPO update

    # eps_cilp:float -> PPO 算法中的 epsilon 值，用于限制 policy 更新的幅度
    eps_clip = 0.2          # clip parameter for PPO
    # gamma:float -> 奖励的折扣系数
    gamma = 0.99            # discount factor

    # lr_actor: float -> actor network 的学习率
    lr_actor = 0.0003       # learning rate for actor network
    # lr_critic: float -> 评论家 network 的学习率
    lr_critic = 0.001       # learning rate for critic network

    # random_seed:int -> 随机数种子
    random_seed = 0         # set random seed if required (0 = no random seed)
    #####################################################

    print("training environment name : " + env_name)

    # env: gym 中的环境类
    env = gym.make(env_name)

    # state space dimension
    # state_dim:int -> 状态空间的数量 如果是 机器狗项目的话，有 22 个 State dim
    state_dim = env.observation_space.shape[0]

    # action space dimension
    if has_continuous_action_space:
        action_dim = env.action_space.shape[0]       # action_dim: int -> 对应动作空间的数量
    else:
        action_dim = env.action_space.n

    ###################### logging ######################

    #### log files for multiple runs are NOT overwritten
    log_dir = "PPO_logs"    # 下面这部分是创建 本次训练的 log 文件夹
    if not os.path.exists(log_dir):
          os.makedirs(log_dir)

    log_dir = log_dir + '/' + env_name + '/'
    if not os.path.exists(log_dir):
          os.makedirs(log_dir)

    #### get number of log files in log directory
    run_num = 0            # 这部分是判断本次训练 log 文件夹下有几个文件，本次训练继续加 1
    current_num_files = next(os.walk(log_dir))[2]
    run_num = len(current_num_files)

    #### create new log file for each run
    log_f_name = log_dir + '/PPO_' + env_name + "_log_" + str(run_num) + ".csv"          # 生成本次训练使用的 log 文件 csv

    print("current logging run number for " + env_name + " : ", run_num)
    print("logging at : " + log_f_name)
    #####################################################

    ################### checkpointing ###################
    run_num_pretrained = 0      #### change this to prevent overwriting weights in same env_name folder   # run_num_pretrained:int -> 这个是用来控制多少个循环会将网络的权重保存到本地的

    directory = "PPO_preTrained"
    if not os.path.exists(directory):
          os.makedirs(directory)

    directory = directory + '/' + env_name + '/'
    if not os.path.exists(directory):
          os.makedirs(directory)


    checkpoint_path = directory + "PPO_{}_{}_{}.pth".format(env_name, random_seed, run_num_pretrained)
    print("save checkpoint path : " + checkpoint_path)
    #####################################################


    ############# print all hyperparameters #############
    print("--------------------------------------------------------------------------------------------")
    print("max training timesteps : ", max_training_timesteps)
    print("max timesteps per episode : ", max_ep_len)
    print("model saving frequency : " + str(save_model_freq) + " timesteps")
    print("log frequency : " + str(log_freq) + " timesteps")
    print("printing average reward over episodes in last : " + str(print_freq) + " timesteps")
    print("--------------------------------------------------------------------------------------------")
    print("state space dimension : ", state_dim)
    print("action space dimension : ", action_dim)
    print("--------------------------------------------------------------------------------------------")
    if has_continuous_action_space:
        print("Initializing a continuous action space policy")
        print("--------------------------------------------------------------------------------------------")
        print("starting std of action distribution : ", action_std)
        print("decay rate of std of action distribution : ", action_std_decay_rate)
        print("minimum std of action distribution : ", min_action_std)
        print("decay frequency of std of action distribution : " + str(action_std_decay_freq) + " timesteps")
    else:
        print("Initializing a discrete action space policy")
    print("--------------------------------------------------------------------------------------------")
    print("PPO update frequency : " + str(update_timestep) + " timesteps")
    print("PPO K epochs : ", K_epochs)
    print("PPO epsilon clip : ", eps_clip)
    print("discount factor (gamma) : ", gamma)
    print("--------------------------------------------------------------------------------------------")
    print("optimizer learning rate actor : ", lr_actor)
    print("optimizer learning rate critic : ", lr_critic)
    if random_seed:
        print("--------------------------------------------------------------------------------------------")
        print("setting random seed to ", random_seed)
        torch.manual_seed(random_seed)
        env.seed(random_seed)
        np.random.seed(random_seed)
    #####################################################

    print("============================================================================================")

    ################# training procedure ################

    # initialize a PPO agent
    ppo_agent = PPO(state_dim, action_dim, lr_actor, lr_critic, gamma, K_epochs, eps_clip, has_continuous_action_space, action_std)
    # ppo_agent: 简化版本的 PPO 模型，里面主要包括 policy_net 和 old_policy_net，每个 policy_net 中包括一个 critic_net 和一个 actor_net，还有一个循环储存数据的 buffer
    # track total training time
    start_time = datetime.now().replace(microsecond=0)
    print("Started training at (GMT) : ", start_time)

    print("============================================================================================")

    # logging file
    log_f = open(log_f_name,"w+")
    log_f.write('episode,timestep,reward\n')

    # printing and logging variables
    print_running_reward = 0                 # print_running_reward:int -> 要打印的 reward
    print_running_episodes = 0               # print_running_episodes:int -> 要打印的 episode

    log_running_reward = 0                   # log_running_reward:int -> 要记录的 reward
    log_running_episodes = 0                 # log_running_episodes:int -> 要记录的 episode

    time_step = 0                            # time_stp: int ->  当前的时间步数
    i_episode = 0                            # i_episode: int -> 当前的 episode 数

    # training loop
    while time_step <= max_training_timesteps:       # 一共训练 max_training_timesteps 个时间步

        state = env.reset()                  # state: np.ndarray -> 环境的状态空间的初始状态
        current_ep_reward = 0                # current_ep_reward:float -> 记录当前 episode 的 reward

        for t in range(1, max_ep_len+1):     # 每个训练的循环迭代中，一共跑 max_ep_len 个 timesteps 数
            #! TODO 这里需要到时候看一下 state 和 action 的维度还有含义
            # select action with policy
            action = ppo_agent.select_action(state)     # 先在 ppo_agent 中选择一个动作
            state, reward, done, _ = env.step(action)   # 再在 env 中计算 下一个时刻的 state，reward，done, info 等信息

            # saving reward and is_terminals
            ppo_agent.buffer.rewards.append(reward)     # 在缓存数据中加入 reward
            ppo_agent.buffer.is_terminals.append(done)  # 在缓存数据中加入是否终止

            time_step +=1                               # time_step += 1
            current_ep_reward += reward                 # 本论迭代的 reward += 本次的 reward

            # update PPO agent
            if time_step % update_timestep == 0:        # 每 update_timestep 个时间步更新一次 PPO agent
                ppo_agent.update()                      # 更新 ppo 内部的网络梯度

            # if continuous action space; then decay action std of ouput action distribution
            """
            如果是连续动作空间中，输出的是动作分布的标准差需要衰减，这里的衰减频率是 action_std_decay_freq 个时间步长，
            每 action_std_decay_freq 个时间步长，就衰减一次，直到 action_std 低于 min_action_std 为止。
            目的主要是：
                1. 平衡前期的探索和利用
                2. 提高策略的稳定性
            """
            if has_continuous_action_space and time_step % action_std_decay_freq == 0:
                ppo_agent.decay_action_std(action_std_decay_rate, min_action_std)

            # log in logging file
            """
            下面是储存到 logging 文件中
            log_freq 和 print_freq 都是每多少个时间步长就记录一次平均奖励，
            保存在文件或者打印出来
            """
            if time_step % log_freq == 0:

                # log average reward till last episode
                log_avg_reward = log_running_reward / log_running_episodes
                log_avg_reward = round(log_avg_reward, 4)

                log_f.write('{},{},{}\n'.format(i_episode, time_step, log_avg_reward))
                log_f.flush()

                log_running_reward = 0
                log_running_episodes = 0

            # printing average reward
            if time_step % print_freq == 0:

                # print average reward till last episode
                print_avg_reward = print_running_reward / print_running_episodes
                print_avg_reward = round(print_avg_reward, 2)

                print("Episode : {} \t\t Timestep : {} \t\t Average Reward : {}".format(i_episode, time_step, print_avg_reward))

                print_running_reward = 0
                print_running_episodes = 0

            # save model weights
            if time_step % save_model_freq == 0:
                print("--------------------------------------------------------------------------------------------")
                print("saving model at : " + checkpoint_path)
                ppo_agent.save(checkpoint_path)
                print("model saved")
                print("Elapsed Time  : ", datetime.now().replace(microsecond=0) - start_time)
                print("--------------------------------------------------------------------------------------------")

            # break; if the episode is over
            if done:
                break

        print_running_reward += current_ep_reward
        print_running_episodes += 1

        log_running_reward += current_ep_reward
        log_running_episodes += 1

        i_episode += 1

    log_f.close()
    env.close()

    # print total training time
    print("============================================================================================")
    end_time = datetime.now().replace(microsecond=0)
    print("Started training at (GMT) : ", start_time)
    print("Finished training at (GMT) : ", end_time)
    print("Total training time  : ", end_time - start_time)
    print("============================================================================================")


if __name__ == '__main__':

    train()
    
    
    
    
    
    
    
